import time

import PIL
import numpy
import numpy as np
from PIL import Image
from torch import nn
from torch.nn import functional as F
import torch
import cv2
from torchvision.transforms import transforms


def get_kernel(kernel_len=16, nsig=10):  # nsig 标准差 ，kernlen=16核尺寸
    GaussianKernel = cv2.getGaussianKernel(kernel_len, nsig) \
                     * cv2.getGaussianKernel(kernel_len, nsig).T
    return GaussianKernel


class Gaussian_kernel(nn.Module):
    def __init__(self,
                 # device,
                 kernel_len, nsig=20):
        super(Gaussian_kernel, self).__init__()
        self.kernel_len = kernel_len
        kernel = get_kernel(kernel_len=kernel_len, nsig=nsig)  # 获得高斯卷积核
        kernel = torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0)  # 扩展两个维度
        # self.weight = nn.Parameter(data=kernel, requires_grad=False).to(device)
        self.weight = nn.Parameter(data=kernel, requires_grad=False)

        self.padding = torch.nn.ReplicationPad2d(int(self.kernel_len/2))

    def forward(self, x):  # x1是用来计算attention的，x2是用来计算的Cs
        x = self.padding(x)
        # 对三个channel分别做卷积
        res = []
        for i in range(x.shape[1]):
            res.append(F.conv2d(x[:, i, :, :], self.weight))
        x_output = torch.cat(res, dim=0)
        return x_output


if __name__ == '__main__':
    ori_image = cv2.imread("t2_opti_debug_1.png")

    kernel = Gaussian_kernel(11, 30).cuda()
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])

    t1 = time.perf_counter()
    ori_image = cv2.resize(ori_image, (224, 224)).astype(np.float32)
    image_cuda = transform(ori_image).unsqueeze(0).cuda()

    out = numpy.transpose(kernel(image_cuda).cpu().numpy(), (1, 2, 0))
    if out.shape != (224, 224, 3):
        out = out[0: 224, 0: 224]
    exposure = cv2.addWeighted(ori_image, 4, out, -4, 128)
    exposure = numpy.clip(exposure, 0, 255).astype(np.uint8)
    exposure = cv2.cvtColor(exposure, cv2.COLOR_BGR2RGB)
    t2 = time.perf_counter()

    print("gpu cost", t2 - t1, " s")
    cv2.imwrite("test_gpu.png", exposure)

